from eegGAN import GANTrainer
from eegGAN.data import MyDataset
from eegGAN.model import UNetGenerator1D, DualPathDiscriminator
import torch
from torch.utils.data import DataLoader
import numpy as np


def run():
    # 配置参数
    config = {
        'noise_dim': 128,
        'g_lr': 0.0002,
        'd_lr': 0.0002,
        'betas': (0.5, 0.999),
        'bis_loss_weight': 1.0,
        'val_interval': 5,
        'save_interval': 10,
        'save_dir': 'D:/Data/FWL/pycharm-workspace/EEG_generate_continue/weights'
    }

    # 设备设置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # 创建模型
    generator = UNetGenerator1D(
        noise_dim=config['noise_dim'],
        bis_dim=1,
        output_dim=256,
        hidden_dim=64
    ).to(device)

    discriminator = DualPathDiscriminator(
        input_dim=256,
        hidden_dim=64,
        bis_prediction=True
    ).to(device)

    # 创建训练器
    trainer = GANTrainer(generator, discriminator, device, config)

    # 创建数据加载器
    dataset = MyDataset("D:/Data/pycharm-workspace/EEG_pytorch/dataset_part/test_data.csv")
    print("datatype size = {}".format(len(dataset)))
    train_dataset = dataset
    val_dataset = dataset
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=128, shuffle=False)

    # 开始训练
    trainer.train(train_loader, val_loader, num_epochs=5000)

    # 绘制训练历史
    trainer.plot_training_history('training_history.png')

    # 生成样本
    samples = trainer.generate_samples(5)
    print(f"生成的样本形状: {samples.shape}")


if __name__ == '__main__':
    run()
